import { useStorage } from '@vueuse/core'
import {
  Ollama,
  type AbortableAsyncIterator,
  type ErrorResponse,
  type ModelResponse,
  type ProgressResponse,
} from 'ollama'

import { cloneDeep, get } from 'lodash-es'
export interface ModelList extends ModelResponse {
  /**状态 */
  status: 'success' | 'error' | 'cancel' | 'loading'
  /**状态信息 */
  msg: string
  /**已下载 */
  completed: number
  /**下载速度 */
  speed: number
  /**当前时间 */
  time: number
  /**取消器 */
  aborter?: AbortableAsyncIterator<ProgressResponse>
}

const emptyObj: ModelList = {
  name: '',
  model: '',
  size: 0,
  modified_at: new Date(),
  digest: '',
  expires_at: new Date(),
  size_vram: 0,
  details: {
    parent_model: '',
    format: '',
    family: '',
    families: [],
    parameter_size: '',
    quantization_level: '',
  },
  status: 'loading',
  msg: '已下载',
  completed: 0,
  time: 0,
  speed: 0,
  aborter: undefined,
}

export const useModelStore = defineStore('model', () => {
  /**本地ollama服务地址 */
  const host = useStorage('ollama-host', 'http://localhost:11434')
  const changeHost = () => {
    if (
      !/^(https?:\/\/)?([a-zA-Z0-9.-]+\.[a-zA-Z]{2,}|localhost)(:\d{1,5})?([/][^\s]*)?$/.test(
        host.value,
      )
    ) {
      ElMessage.error(
        '请输入正确的ollama服务地址，必须以http://或https://开头，不能携带任何query参数',
      )
      return
    }
    ollama.value = new Ollama({ host: host.value })
    list.value = []
    cache.value.clear()
    getList()
    ElMessage.success('已切换ollama服务地址')
  }

  /**ollama服务实例 */
  const ollama = shallowRef(new Ollama({ host: host.value }))
  /**模型列表 */
  const list = ref<ModelList[]>([])
  const cache = useStorage<Set<string>>('pull-cache', new Set())

  /**获取模型列表 */
  const getList = async () => {
    try {
      const { models = [] } = await ollama.value.list()
      list.value = models.map((item) => ({
        ...item,
        status: 'success',
        completed: item.size,
        msg: '已下载',
        time: 0,
        speed: 0,
        aborter: undefined,
      }))

      if (cache.value.size) {
        const arr = Array.from(cache.value)
        list.value.push(
          ...arr.map((name) => {
            const item: ModelList = { ...cloneDeep(emptyObj), name, model: name, status: 'cancel' }
            return item
          }),
        )
      }
    } catch (error) {
      console.error('get ollama model list error:', error)
      ElMessage.error('获取模型列表失败，请检查ollama服务是否已启动')
    }
  }
  /**模型列表相关操作 */

  // 下载（更新）模型
  const pullModel = async (name: string) => {
    if (!name) {
      ElMessage.error('请输入模型名称')
      return
    }
    // 不存在模型（包括缓存）则添加一条记录
    if (!list.value.some((item) => item.name === name)) {
      list.value.push({ ...cloneDeep(emptyObj), name, model: name })
      cache.value.add(name)
    }
    // 重新查找对应记录
    const target = list.value.find((item) => item.name === name)!
    const res = await ollama.value.pull({ model: name, stream: true })
    // 添加取消器
    target.aborter = res
    try {
      for await (const element of res) {
        const { status, total = 0, completed = 0 } = element
        target.size = total
        target.msg = status
        target.status = status === 'success' ? 'success' : 'loading'
        const time = Date.now()
        if (time - target.time > 1000) {
          target.speed = !completed ? 0 : completed - target.completed
          target.completed = completed
          target.time = time
        }
      }
    } catch (error) {
      const { message } = error as Error
      if (message.includes('aborted')) {
        target.status = 'cancel'
        target.msg = '已取消'
      } else {
        target.status = 'error'
        target.msg = message
      }
    }
    target.aborter = undefined
  }
  // 删除模型
  const delModel = async (name: string) => {
    if (cache.value.has(name)) {
      cache.value.delete(name)
      ElMessage.success(`删除成功:${name}，已下载的缓存将在重启电脑后被删除。或选择手动删除。`)
    } else {
      try {
        const { status } = await ollama.value.delete({ model: name })
        ElMessage.success(`删除成功:${status}`)
      } catch (err) {
        const { error } = err as ErrorResponse
        ElMessage.error(`删除失败:${error}`)
      }
    }
    getList()
  }
  return {
    host,
    changeHost,
    ollama,
    list,
    cache,
    getList,
    pullModel,
    delModel,
  }
})
